#include <algorithm>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <thread>
#include <immintrin.h>
#include "objs/dist_ispc.h"

#include "CycleTimer.h"

using namespace std;

static constexpr int nThread = 8;

typedef struct {
  // Control work assignments
  int start, end;
  int Cur_start, Cur_end;

  // Shared by all functions
  double *data;
  double *clusterCentroids;
  int *clusterAssignments;
  double *currCost;
  int M, N, K;
  double distTime;
  double initialTime;
  double judgeTime;
  double loopTime;
  double deleteTime;
} WorkerArgs;

struct DistWorkerArgs{
  int start_index, end_index;
  double local_dist;
};

/**
 * Checks if the algorithm has converged.
 * 
 * @param prevCost Pointer to the K dimensional array containing cluster costs 
 *    from the previous iteration.
 * @param currCost Pointer to the K dimensional array containing cluster costs 
 *    from the current iteration.
 * @param epsilon Predefined hyperparameter which is used to determine when
 *    the algorithm has converged.
 * @param K The number of clusters.
 * 
 * NOTE: DO NOT MODIFY THIS FUNCTION!!!
 */
static bool stoppingConditionMet(double *prevCost, double *currCost,
                                 double epsilon, int K) {
  for (int k = 0; k < K; k++) {
    if (abs(prevCost[k] - currCost[k]) > epsilon)
      return false;
  }
  return true;
}

/**
 * Computes L2 distance between two points of dimension nDim.
 * 
 * @param x Pointer to the beginning of the array representing the first
 *     data point.
 * @param y Poitner to the beginning of the array representing the second
 *     data point.
 * @param nDim The dimensionality (number of elements) in each data point
 *     (must be the same for x and y).
 */
inline double dist(double *x, double *y, int nDim) {
  double accum = 0.0;
  for (int i = 0; i < nDim; i++) {
    accum += pow((x[i] - y[i]), 2);
  }
  return sqrt(accum);
}

// inline double dist_instrinsic(double* x, double *y, int nDim)
// {
//   double accum = 0.0f;

//   __m512i x_vec = _mm512_load_epi32(x);
//   __m512i y_vec = _mm512_load_epi32(y);


// }

void dist_work(DistWorkerArgs& args, double* x, double* y)
{
  for(int index = args.start_index; index < args.end_index; index ++)
  {
    args.local_dist += pow((x[index] - y[index]), 2);
  }
}

double dist_multithread(double*x ,double* y, int nDim)
{
  DistWorkerArgs args[nThread];
  std::thread workers[nThread];

  double accum = 0.0f;
  for(int i = 0; i < nThread; i ++)
  {
    args[i].start_index = i * nDim / nThread;
    args[i].end_index   = std::min((i + 1) * nDim / nThread, nDim);
    args[i].local_dist  = 0.0f;
  }

  for(int i = 1; i < nThread; i ++)
  {
    workers[i] = std::thread(dist_work, std::ref(args[i]), x, y);
  }
  
  dist_work(args[0], x, y);

  for(int i = 1; i < nThread; i ++)
  {
    workers[i].join();
  }

  for(int i = 0; i < nThread; i ++)
  {
    accum += args[i].local_dist;
  }

  return sqrt(accum);
}

/**
 * Assigns each data point to its "closest" cluster centroid.
 */
void computeAssignments(WorkerArgs *const args) {
  double *minDist = new double[args->M];
  
  double startTime = CycleTimer::currentSeconds();
  // Initialize arrays
  for (int m =0; m < args->M; m++) {
    minDist[m] = 1e30;
    args->clusterAssignments[m] = -1;
  }
  double endTime = CycleTimer::currentSeconds();
  // printf("Initialization time: %.3f ms\n", (endTime - startTime) * 1000);

  args->initialTime += (endTime - startTime) * 1000;

  startTime = CycleTimer::currentSeconds();
  // Assign datapoints to closest centroids
  for (int k = args->start; k < args->end; k++) {
    for (int m = 0; m < args->M; m++) {
      double startTime = CycleTimer::currentSeconds();
      // double d = dist(&args->data[m * args->N],
      //                 &args->clusterCentroids[k * args->N], args->N);
      // double d = dist_multithread(
      //   &args->data[m * args->N],
      //   &args->clusterCentroids[k * args->N], args->N
      // );
      double d = 0.0f;
      d = ispc::dist_ispc(
        &args->data[m * args->N],
        &args->clusterCentroids[k * args->N],
        args->N
      );
      double endTime   = CycleTimer::currentSeconds();
      args->distTime   += (endTime - startTime) * 1000;
      startTime = CycleTimer::currentSeconds();
      if (d < minDist[m]) {
        minDist[m] = d;
        args->clusterAssignments[m] = k;
      }
      endTime = CycleTimer::currentSeconds();
      args->judgeTime += (endTime - startTime) * 1000;
    }
  }
  endTime = CycleTimer::currentSeconds();
  args->loopTime += (endTime - startTime) * 1000;


  startTime = CycleTimer::currentSeconds();
  free(minDist);
  endTime = CycleTimer::currentSeconds();
  args->deleteTime += (endTime - startTime) * 1000;
}

void worker_start(WorkerArgs& args)
{
  int element_count = args.Cur_end - args.Cur_start;
  double *minDist   = new double[element_count + 2];
  double startTime = 0.0f;
  double endTime   = 0.0f;

  // Initialize arrays
  startTime = CycleTimer::currentSeconds();
  for(int i = 0; i < element_count; i++) {
    minDist[i] = 1e30;
    args.clusterAssignments[args.Cur_start + i] = -1;
  }
  endTime = CycleTimer::currentSeconds();
  // printf("Initialization time: %.3f ms\n", (endTime - startTime) * 1000);
  args.initialTime = (endTime - startTime) * 1000;

  startTime = 0.0f;
  endTime   = 0.0f;
  // Assign datapoints to closest centroids
  for(int k = args.start; k < args.end; k ++)
  {
    for(int m = args.Cur_start; m < args.Cur_end; m ++)
    {
      startTime = CycleTimer::currentSeconds();
      // double d  = dist(
      //   &args.data[args.N * m],
      //   &args.clusterCentroids[args.N * k],
      //   args.N
      // );
      double d = 0.0f;
      d = ispc::dist_ispc(
        &args.data[args.N * m],
        &args.clusterCentroids[args.N * k],
        args.N
      );
      // double d = dist_multithread(
      //   &args.data[args.N * m],
      //   &args.clusterCentroids[args.N * k],
      //   args.N
      // );
      endTime   = CycleTimer::currentSeconds();
      args.distTime += (endTime - startTime) * 1000;

      if(d < minDist[m - args.Cur_start]) 
      {
        minDist[m - args.Cur_start] = d;
        args.clusterAssignments[m] = k;
      }
    }
  }

  delete[] minDist;
}

// MultiThread Assignment
void Compute_MultiThreadAssignment(WorkerArgs& out_args)
{
  std::thread workers[nThread];
  WorkerArgs args[nThread];

  for(int i = 0; i < nThread; i ++)
  {
    args[i].Cur_start = i * out_args.M / nThread;
    args[i].Cur_end   = std::min((i + 1) * out_args.M / nThread, out_args.M);
    
    args[i].start              = out_args.start;
    args[i].end                = out_args.end;
    args[i].data               = out_args.data;
    args[i].clusterCentroids   = out_args.clusterCentroids;
    args[i].clusterAssignments = out_args.clusterAssignments;
    args[i].M                  = out_args.M;
    args[i].N                  = out_args.N;
    args[i].K                  = out_args.K;
    args[i].distTime           = 0.0;
    args[i].initialTime        = 0.0f;
  }

  for(int i = 1; i < nThread; i ++)
  {
    workers[i] = std::thread(worker_start, std::ref(args[i]));
  }

  worker_start(std::ref(args[0]));

  for(int i = 1; i < nThread; i ++)
  {
    workers[i].join();
  }

  double distTime = 0.0;
  double initialTime = 0.0;
  for(int i = 0; i < nThread; i ++)
  {
    distTime += args[i].distTime;
    initialTime += args[i].initialTime;
  }
  out_args.distTime = distTime;
  out_args.initialTime = initialTime;

}

/**
 * Given the cluster assignments, computes the new centroid locations for
 * each cluster.
 */
void computeCentroids(WorkerArgs *const args) {
  int *counts = new int[args->K];

  // Zero things out
  for (int k = 0; k < args->K; k++) {
    counts[k] = 0;
    for (int n = 0; n < args->N; n++) {
      args->clusterCentroids[k * args->N + n] = 0.0;
    }
  }


  // Sum up contributions from assigned examples
  for (int m = 0; m < args->M; m++) {
    int k = args->clusterAssignments[m];
    for (int n = 0; n < args->N; n++) {
      args->clusterCentroids[k * args->N + n] +=
          args->data[m * args->N + n];
    }
    counts[k]++;
  }

  // Compute means
  for (int k = 0; k < args->K; k++) {
    counts[k] = max(counts[k], 1); // prevent divide by 0
    for (int n = 0; n < args->N; n++) {
      args->clusterCentroids[k * args->N + n] /= counts[k];
    }
  }

  free(counts);
}

/**
 * Computes the per-cluster cost. Used to check if the algorithm has converged.
 */
void computeCost(WorkerArgs *const args) {
  double *accum = new double[args->K];

  // Zero things out
  for (int k = 0; k < args->K; k++) {
    accum[k] = 0.0;
  }

  // Sum cost for all data points assigned to centroid
  for (int m = 0; m < args->M; m++) {
    int k = args->clusterAssignments[m];
    accum[k] += dist(&args->data[m * args->N],
                     &args->clusterCentroids[k * args->N], args->N);
  }

  // Update costs
  for (int k = args->start; k < args->end; k++) {
    args->currCost[k] = accum[k];
  }

  free(accum);
}

/**
 * Computes the K-Means algorithm, using std::thread to parallelize the work.
 *
 * @param data Pointer to an array of length M*N representing the M different N 
 *     dimensional data points clustered. The data is layed out in a "data point
 *     major" format, so that data[i*N] is the start of the i'th data point in 
 *     the array. The N values of the i'th datapoint are the N values in the 
 *     range data[i*N] to data[(i+1) * N].
 * @param clusterCentroids Pointer to an array of length K*N representing the K 
 *     different N dimensional cluster centroids. The data is laid out in
 *     the same way as explained above for data.
 * @param clusterAssignments Pointer to an array of length M representing the
 *     cluster assignments of each data point, where clusterAssignments[i] = j
 *     indicates that data point i is closest to cluster centroid j.
 * @param M The number of data points to cluster.
 * @param N The dimensionality of the data points.
 * @param K The number of cluster centroids.
 * @param epsilon The algorithm is said to have converged when
 *     |currCost[i] - prevCost[i]| < epsilon for all i where i = 0, 1, ..., K-1
 */
void kMeansThread(double *data, double *clusterCentroids, int *clusterAssignments,
               int M, int N, int K, double epsilon, bool using_multithread) {

  // Used to track convergence
  double *prevCost = new double[K];
  double *currCost = new double[K];

  // The WorkerArgs array is used to pass inputs to and return output from
  // functions.
  WorkerArgs args;
  args.data = data;
  args.clusterCentroids = clusterCentroids;
  args.clusterAssignments = clusterAssignments;
  args.currCost = currCost;
  args.M = M;
  args.N = N;
  args.K = K;
  args.distTime  = 0.0f;
  args.judgeTime = 0.0f;
  args.loopTime  = 0.0f;

  // Initialize arrays to track cost
  for (int k = 0; k < K; k++) {
    prevCost[k] = 1e30;
    currCost[k] = 0.0;
  }

  double AssignmentTime = 0.0f;
  double CentroidTime   = 0.0f;
  double CostTime       = 0.0f;
  double startTime      = 0.0f;
  double endTime        = 0.0f;

  bool using_thread = using_multithread;

  /* Main K-Means Algorithm Loop */
  int iter = 0;
  while (!stoppingConditionMet(prevCost, currCost, epsilon, K)) {
    // Update cost arrays (for checking convergence criteria)
    for (int k = 0; k < K; k++) {
      prevCost[k] = currCost[k];
    }

    // Setup args struct
    args.start = 0;
    args.end = K;

    startTime = CycleTimer::currentSeconds();
    // computeAssignments(&args);
    // Compute_MultiThreadAssignment(args);
    if(using_thread) Compute_MultiThreadAssignment(args);
    else             computeAssignments(&args);
    endTime = CycleTimer::currentSeconds();
    AssignmentTime += (endTime - startTime) * 1000;

    startTime = CycleTimer::currentSeconds();
    computeCentroids(&args);
    endTime = CycleTimer::currentSeconds();
    CentroidTime += (endTime - startTime) * 1000;

    startTime = CycleTimer::currentSeconds();
    computeCost(&args);
    endTime = CycleTimer::currentSeconds();
    CostTime += (endTime - startTime) * 1000;

    iter++;
  }

  printf("AssignmentTime: %.3f ms\n", AssignmentTime);
  printf("In Assignment, Dist Sum Time: %.3f ms\n", args.distTime);
  printf("In Assignment, Initialize Time : %.3f ms\n", args.initialTime);
  printf("In Assignment, Judge Time : %.3f ms\n", args.judgeTime);
  printf("In Assignment, Loop Time : %.3f ms\n", args.loopTime);
  printf("In Assginment, deletetime : %.3f ms\n", args.deleteTime);
  if(using_thread)
    {
      printf("In Assignment, Using %d threads, each thread evarage time : %.3f ms\n", nThread, args.distTime / nThread);
      printf("In Assignment, Using %d threads, initialize using time : %.3f ms\n", nThread, args.initialTime);
    }
  printf("CentroidTime: %.3f ms\n", CentroidTime);
  printf("CostTime: %.3f ms\n", CostTime);

  free(currCost);
  free(prevCost);
}
